from flask import Flask, request, jsonify
import torch
import numpy as np
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import spectral
import os
import time
from utils import HybridSN, applyPCA, padWithZeros

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
matplotlib.use('TkAgg')
import scipy.io as sio
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix, accuracy_score, classification_report, cohen_kappa_score
import spectral
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pickle
import pymysql
import scipy.io as sio
from sklearn.decomposition import PCA
import numpy as np
# from flask import Flask, render_template, request, redirect, url_for, flash
from flask import Flask, request, jsonify, render_template, redirect, url_for, flash
import scipy.io as sio


class_num = 16
app = Flask(__name__)
app.secret_key = 'your_secret_key_here'  # 设置一个密钥用于Flask会话

# 加载模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = HybridSN().to(device)
model.load_state_dict(torch.load(os.path.join('model1.pth'),weights_only=True))
model.eval()

# 配置文件上传的保存路径
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route('/', methods=["GET","POST"])
def classify():
    if request.method == 'GET':
        return render_template('index.html')
    elif request.method == 'POST':
        file1 = request.files['image1']
        file2 = request.files['image2']

        if file1.filename == '' or file2.filename == '':
            flash("未选择文件或文件名为空")
            return redirect(url_for('classify'))

        file1_path = os.path.join(app.config['UPLOAD_FOLDER'], file1.filename)
        file2_path = os.path.join(app.config['UPLOAD_FOLDER'], file2.filename)
        file1.save(file1_path)
        file2.save(file2_path)

        flash("文件已成功上传并保存到本地！")
        print("文件已成功上传并保存到本地！")

        # X
        mat_data = sio.loadmat(file1.filename)
        print(mat_data.keys())
        print(mat_data.values())
        # X = None
        for key in mat_data.keys():
            print("key",key)
            value = mat_data[key]
            try:
                if len(value.shape) == 3:
                    print("key_X", key)
                    X = sio.loadmat(file1.filename)[key]
            except AttributeError:
                continue
        print(X.shape)

        # y
        mat_data1 = sio.loadmat(file2.filename)
        print(mat_data1.keys())
        y = None
        for key in mat_data1.keys():
            print(key)
            value = mat_data1[key]
            try:
                if len(value.shape) == 2:
                    print("key_y",key)
                    y = sio.loadmat(file2.filename)[key]
            except AttributeError:
                continue
        print(y.shape)


        height = y.shape[0]
        width = y.shape[1]

        patch_size = 25
        pca_components = 30

        X = applyPCA(X, numComponents=pca_components)
        X = padWithZeros(X, patch_size // 2)

        # 逐像素预测类别
        outputs = np.zeros((height, width))
        for i in range(height):
            for j in range(width):
                if int(y[i, j]) == 0:
                    continue
                else:
                    image_patch = X[i:i + patch_size, j:j + patch_size, :]
                    image_patch = image_patch.reshape(1, image_patch.shape[0], image_patch.shape[1],
                                                      image_patch.shape[2],
                                                      1)
                    X_test_image = torch.FloatTensor(image_patch.transpose(0, 4, 3, 1, 2)).to(device)
                    prediction = model(X_test_image)
                    prediction = np.argmax(prediction.detach().cpu().numpy(), axis=1)
                    outputs[i][j] = prediction + 1
                if i % 20 == 0:
                    print('... ... row ', i, ' handling ... ...')
        print(outputs.shape)
        print(outputs.dtype)
        print(outputs.min(), outputs.max())
        predict_image = spectral.imshow(classes=outputs.astype(int), figsize=(5, 5))
        print(predict_image)

        # 生成新的图像文件
        output_image_path = os.path.join('static', f'output_image_{int(time.time())}.png')
        plt.savefig(output_image_path)
        plt.close()

        return jsonify({'image_url': output_image_path})
    return render_template('index.html')

    #     plt.savefig('static/output_image.png')
    #     return jsonify({'image_url': '/static/output_image.png'})
    #     # plt.savefig('output_image.png')
    # return render_template('index.html')
            # # SalinasA
            # mat_data = sio.loadmat('SalinasA_gt.mat')
            # print(mat_data.keys())
            # X = sio.loadmat('SalinasA_corrected.mat')['salinasA_corrected']
            # print(X.shape)  # 打印数据X的维度信息
            # y = sio.loadmat('SalinasA_gt.mat')['salinasA_gt']

            # Indian_pines
            # mat_data = sio.loadmat('Indian_pines_corrected .mat')
            # print(mat_data.keys())
            # X = sio.loadmat('Indian_pines_corrected .mat')['indian_pines_corrected']
            # print(X.shape)  # 打印数据X的维度信息
            # y = sio.loadmat('Indian_pines_gt.mat')['indian_pines_gt']

@app.route('/static/<path:path>')
def serve_static(path):
    return app.send_static_file(path)


if __name__ == '__main__':
    app.run(debug=True)
